import os
import zipfile
from io import BytesIO

import requests
import torch
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms, datasets


# 设置标签文字与numpy的转换


class CustomImageDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        """
        Args:
            root_dir (str): 图片根目录，要求每个子文件夹名为类别名
            transform (callable, optional): 对图片进行处理的 transform
        """
        self.root_dir = root_dir
        self.transform = transform

        # 使用 ImageFolder 自动处理文件夹结构和标签
        self.dataset = datasets.ImageFolder(root=root_dir, transform=self.transform)

    def __getitem__(self, idx):
        # 返回 (image, label)
        image, label = self.dataset[idx]
        return image, label

    def __len__(self):
        return len(self.dataset)
